#!/usr/bin/env python3

# Copyright © 2025 Wenze Wei
#
# This file is part of Pisces L1.
#
# Licensed under the Creative Commons Attribution-NonCommercial 4.0 International License (CC BY-NC 4.0).
# You may not use this file except in compliance with the License.
# Commercial use is strictly prohibited.
# You may obtain a copy of the License at
#
#     https://creativecommons.org/licenses/by-nc/4.0/
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
import torch
import torch.nn as nn
from typing import Optional

class YaRNRotaryEmbedding(nn.Module):
    """Implementation of iRoPE (Interpolation Rotary Position Encoding) based on the Llama 4 Scout architecture.
    Supports a 10M context window and uses interpolated rotary position encoding with staggered attention mechanism.
    """
    def __init__(self,
                 dim: int,
                 max_position_embeddings: int = 10485760,
                 base: int = 1000000,
                 scale: float = 1.0,
                 original_max_position_embeddings: int = 10485760,
                 device: Optional[torch.device] = None):
        """Initialize the YaRNRotaryEmbedding module.

        Args:
            dim (int): Dimension of the input tensor.
            max_position_embeddings (int, optional): Maximum number of position embeddings. Defaults to 10485760 (10M).
            base (int, optional): Base value for calculating inverse frequencies. Defaults to 1000000.
            scale (float, optional): Scaling factor for iRoPE interpolation. Defaults to 1.0.
            original_max_position_embeddings (int, optional): Original maximum number of position embeddings. Defaults to 65536.
            device (Optional[torch.device], optional): Device to place the tensors on. Defaults to None.
        """
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        self.scale = scale
        self.original_max_position_embeddings = original_max_position_embeddings
        # Calculate and register inverse frequencies for rotary position encoding as buffer
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
    
    def forward(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
        """Forward pass of the YaRNRotaryEmbedding module.

        Args:
            x (torch.Tensor): Input tensor.
            seq_len (int): Length of the sequence.

        Returns:
            torch.Tensor: Tensor with rotary position encoding applied.
        """
        device = x.device
        
        # Validate sequence length bounds
        if seq_len > self.max_position_embeddings:
            print(f"🟧\tSequence length {seq_len} exceeds max {self.max_position_embeddings}, using max length")
            seq_len = self.max_position_embeddings
        if seq_len <= 0:
            print(f"🟧\tInvalid sequence length {seq_len}, using 1")
            seq_len = 1
            
        # Create a tensor representing positions from 0 to seq_len - 1
        t = torch.arange(seq_len, device=device, dtype=torch.float32)
        
        # Validate tensor dimensions
        if x.shape[-1] != self.dim:
            print(f"🟧\tInput dimension {x.shape[-1]} doesn't match expected {self.dim}")
            
        # Apply iRoPE interpolation scaling
        t = t / self.scale if self.scale > 0 else t
        
        # Temperature scaling optimization
        if seq_len > self.original_max_position_embeddings:
            scaling_factor = (self.original_max_position_embeddings / seq_len) ** 0.5
            t = t * scaling_factor
            
        # Ensure frequencies are finite
        if torch.any(~torch.isfinite(self.inv_freq)):
            print("🟧\tNon-finite values in inv_freq detected, using safe values")
            inv_freq = torch.clamp(self.inv_freq, min=1e-6, max=1e6)
        else:
            inv_freq = self.inv_freq
            
        # Compute frequencies by taking the outer product of positions and inverse frequencies
        freqs = torch.outer(t, inv_freq)
        
        # Ensure frequencies are finite
        freqs = torch.clamp(freqs, min=-1e6, max=1e6)
        
        # Compute cosine and sine values of frequencies
        cos = freqs.cos()
        sin = freqs.sin()
        
        return self._rotate_half(x, cos, sin)
    
    @staticmethod
    def _rotate_half(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        """Apply rotary position encoding by rotating half of the input tensor.

        Args:
            x (torch.Tensor): Input tensor.
            cos (torch.Tensor): Cosine values of frequencies.
            sin (torch.Tensor): Sine values of frequencies.

        Returns:
            torch.Tensor: Tensor with rotary position encoding applied.
        """
        # Validate tensor dimensions
        last_dim = x.shape[-1]
        if last_dim % 2 != 0:
            print(f"🟧\tLast dimension {last_dim} is odd, padding to even")
            padding = torch.zeros(*x.shape[:-1], 1, device=x.device, dtype=x.dtype)
            x = torch.cat([x, padding], dim=-1)
            last_dim += 1
            
        # Ensure safe splitting indices
        split_idx = last_dim // 2
        
        # Split the tensor into two halves
        x1 = x[..., :split_idx]
        x2 = x[..., split_idx:]
        
        # Ensure cos and sin have compatible shapes
        if cos.shape[-1] != split_idx:
            cos = cos[..., :split_idx]
            sin = sin[..., :split_idx]
        
        # Apply rotation
        try:
            rotated = torch.cat((-x2 * sin + x1 * cos, x1 * sin + x2 * cos), dim=-1)
        except RuntimeError as e:
            print(f"🟧\tRotation failed: {e}, returning original")
            return x
        
        return rotated
    
    def extra_repr(self) -> str:
        """Return a string representation of the module's parameters.

        Returns:
            str: String containing the module's parameter values.
        """
        return (f"dim={self.dim}, max_position_embeddings={self.max_position_embeddings}, "
                f"base={self.base}, scale={self.scale}, "
                f"original_max_position_embeddings={self.original_max_position_embeddings}")